"""颜色跟踪模块
用于跟踪特定颜色的目标，并控制机器人跟随目标移动
"""

import cv2
import numpy as np
from typing import Dict, Any, Optional, Tuple
from .vision_base import VisionBase
from lab_conf import color_range
import time
import hexapod
import Serial_Servo_Running as SSR
import asyncio


class ColorTracking(VisionBase):
    """颜色跟踪类"""

    def __init__(self, config: Dict[str, Any]):
        """
        初始化颜色跟踪模块

        Args:
            config: 配置字典，包含颜色跟踪参数
        """
        super().__init__(config)
        self.target_color = config.get("target_color", "red")
        self.min_radius = config.get("min_radius", 5)
        self.max_radius = config.get("max_radius", 35)
        self.angle_factor = config.get("angle_factor", 0.167)
        self.turn_threshold = config.get("turn_threshold", 80)
        self.ball_x = 0
        self.radius = 0
        self.Get_Angle = False
        self.angle = 0.0

        # 初始化机器人
        hexapod.camera_pos_init()
        hexapod.hexapod_init()

    def _get_target_position(
        self, frame: np.ndarray
    ) -> Tuple[Optional[Tuple[int, int]], Optional[float]]:
        """
        获取目标位置和半径

        Args:
            frame: 输入图像帧

        Returns:
            目标中心坐标和半径
        """
        # 图像预处理
        frame = cv2.resize(frame, (self.width, self.height))
        frame = cv2.GaussianBlur(frame, (3, 3), 0)
        lab = cv2.cvtColor(frame, cv2.COLOR_BGR2LAB)

        # 颜色分割
        mask = cv2.inRange(
            lab, color_range[self.target_color][0], color_range[self.target_color][1]
        )

        # 形态学处理
        kernel = np.ones((3, 3), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)

        # 查找轮廓
        contours = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[
            -2
        ]
        if len(contours) > 0:
            c = max(contours, key=cv2.contourArea)
            ((x, y), radius) = cv2.minEnclosingCircle(c)
            return (int(x), int(y)), int(radius)
        return None, None

    async def process_frame(self, frame: np.ndarray) -> Dict[str, Any]:
        """
        处理图像帧，跟踪目标

        Args:
            frame: 输入图像帧

        Returns:
            处理结果字典
        """
        # 获取目标位置
        center, radius = self._get_target_position(frame)

        result = {"detected": False, "center": None, "radius": 0, "angle": 0}

        if center is not None:
            x, y = center
            self.ball_x = x
            self.radius = radius

            # 计算转向角度
            if not self.Get_Angle:
                self.angle = int((x - self.width / 2) * self.angle_factor)
                self.Get_Angle = True
            else:
                self.Get_Angle = False

            result.update(
                {
                    "detected": True,
                    "center": center,
                    "radius": radius,
                    "angle": self.angle,
                }
            )

            # 在图像上绘制结果
            cv2.circle(frame, center, radius, (0, 255, 255), 2)
            cv2.circle(frame, center, 5, (0, 255, 255), -1)

        # 添加调试信息
        if self.debug:
            frame = self.add_text_to_image(
                frame, f"Tracking: {self.target_color}", (10, 30)
            )
            frame = self.add_text_to_image(
                frame, f"Detected: {result['detected']}", (10, 60)
            )
            frame = self.draw_crosshair(frame, (self.width // 2, self.height // 2))

        return {"result": result, "frame": frame}

    async def get_vision_data(self) -> Dict[str, Any]:
        """
        获取视觉处理数据

        Returns:
            包含处理结果和时间戳的字典
        """
        frame = await self.get_frame()
        if frame is not None:
            result = await self.process_frame(frame)
            return {
                **result["result"],
                "timestamp": time.time(),
                "frame": result["frame"] if self.debug else None,
            }
        return {
            "detected": False,
            "center": None,
            "radius": 0,
            "angle": 0,
            "timestamp": time.time(),
            "frame": None,
        }

    async def follow_target(self) -> None:
        """控制机器人跟随目标"""
        if self.Get_Angle:
            if self.ball_x < self.width / 2 - self.turn_threshold:
                hexapod.turn(-3, 100)
                self.Get_Angle = False
            elif self.ball_x > self.width / 2 + self.turn_threshold:
                hexapod.turn(3, 100)
                self.Get_Angle = False
            else:
                if self.min_radius < self.radius <= self.max_radius:
                    SSR.run_ActionGroup("41", 1)
                else:
                    await asyncio.sleep(0.01)
        else:
            await asyncio.sleep(0.01)
